perf(kernel): import pyarrow on the first result instead of in connect() (stacked on #1049) - #1053

What type of PR is this?

• Other (performance)

Description

Stacked on #1049 (lazy pandas). Its two commits show in this diff until it merges, and without it pandas would still import numpy in connect().

A kernel connect() imports pyarrow, and with it numpy. client.py, utils.py and result_set.py import pyarrow at module load as an optional dependency. The kernel client imports its result set and type mapping, which need pyarrow, at module load too. Besides the import itself (~70 ms here), numpy starts OpenBLAS's thread pool: one thread per core, busy-waiting for a while after start. A bare import numpy burns 0.33 CPU-seconds in the following idle second on 8 cores (0.00 with OPENBLAS_NUM_THREADS=1). During a 200-thread connect burst that spin was half the process's CPU and took ~65 ms of connect() p50 (measured diagnostically with OPENBLAS_NUM_THREADS=1). With several processes starting at once it competes with every connection being opened.

The three modules now get pyarrow from optional_pyarrow(). It returns the real module if pyarrow is already imported, None if it isn't installed, and otherwise a stand-in that imports pyarrow on first attribute access. That first access also rebinds pyarrow in those modules to the real module, so reads on the query path are plain module attribute lookups (through the stand-in a read costs 63 ns instead of 9 ns, 12 reads per query). Installation is decided with importlib.util.find_spec, so every pyarrow is None check, and the warning client.py logs when pyarrow is missing, behave as before. The module attributes keep their names, and type checkers still see a plain import pyarrow. The kernel client imports its result set and type mapping at the first result and the first bound parameter. The Thrift and SEA backends are unchanged and still import pyarrow when they load. A kernel connection now imports pyarrow and numpy at its first execute(), and pandas (#1049) at its first row conversion. One edge case changes: pyarrow installed but failing to import (a broken install, e.g. a numpy ABI mismatch). Before, every module treated it as absent. Now client, utils and result_set treat it as available, and the first attribute use raises an AttributeError that carries the ImportError, so hasattr and getattr(..., default) still work. The Thrift and SEA backends still import it when they load and treat it as missing, so the Thrift result format is unchanged. In that case the "pyarrow is not installed" warning isn't logged, and connection telemetry reports enable_arrow=True. Making those agree would mean importing pyarrow in connect(), which this change exists to avoid.

How is this tested?

• Unit tests
• Manually

New tests, in fresh interpreters with a stand-in kernel wheel: importing databricks.sql, client, utils, result_set and the kernel client loads none of pyarrow, numpy or pandas. connect(use_kernel=True) doesn't either, and building the first result set does import pyarrow. With pyarrow made unimportable, the three modules' pyarrow is None. The stand-in resolves to the real classes and hands back the real module once pyarrow is imported. The full unit suite, the realkernel sets and mypy pass.

Local mock server adding 20 ms per request, 8 cores, kernel main, connector #1049 + #1051 with and without this change, fresh processes, interleaved:

In the single process, the first result arrives at about the same time or a little later, because pyarrow's import moves there. With four processes starting together, the first query finishes earlier, because numpy's spinning threads no longer overlap the connect burst. The c800 open-time row is measured from process launch, using each session's open timestamp, not c800's own connect() percentiles. c800 starts its threads in a loop, and the import's GIL use delays that loop, so before this change half its threads started connect() 0.3 s or more late and their timings missed that wait.

Steady state, vs #1049 alone (isolated run, c800 4 × 200, 20 ms mock, 60 s, 16 cores, medians of 2 rounds): QPS 5548 → 5582, CPU per query 1.429 → 1.434 ms, each session's open time from process launch p50 1.68 → 1.63 s, all sessions open 1.98 → 1.88 s. Before the rebinding, steady QPS was 2% lower.

Live, on a Reyden warehouse (1 process × 200 threads + 1 query, 8 cores, 5 interleaved runs, a kernel with #453/#455 in both arms), main vs. #1049 + #1051 + this PR together: connect() p50 (from a start barrier) 688 → 368 ms, all sessions open 0.80 → 0.43 s, first result p50 0.83 → 0.77 s, medians of runs.

Live benchmark (staging Large Reyden warehouse)

TPC-H Q6 on a Large Real-Time warehouse, 800 sessions (4 processes x 200 threads), cold start, each session querying back to back for 120 s, 16 client cores. Isolated A/B: only this change vs #1049 c193711 (this PR is stacked on it), 3 interleaved rounds, nothing else on the warehouse. Measured at 517c4c8. Median (min-max):

QPS within noise; small cold-start connect gain.

Related Tickets & Documents

RE-17633 (connection-creation time at high concurrency).

This pull request and its description were written by Isaac.